// Copyright 2024 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import {
  TokenType,
  useAuthState,
  useGetAuthToken,
} from '@/common/components/auth_state_provider';
import { useSingleton } from '@/generic_libs/hooks/singleton';
import { PrpcClient } from '@/generic_libs/tools/prpc_client';
import { getObjectId } from '@/generic_libs/tools/utils';
import { Constructor } from '@/generic_libs/types';

export interface PrpcServiceClientOptions<S, Params extends unknown[] = []> {
  /**
   * The host of the pRPC server.
   */
  readonly host: string;
  /**
   * Additional headers to be passed to the RPC call. This can be used to pass
   * [gRPC metadata](https://grpc.io/docs/guides/metadata/#headers) to the
   * server.
   *
   * Note that this does not override the headers set by `PrpcClient`, including
   *  * accept
   *  * content-type
   *  * authorization
   */
  readonly additionalHeaders?: HeadersInit;
  /**
   * If true, use HTTP instead of HTTPS. Defaults to `false`.
   */
  readonly insecure?: boolean;
  /**
   * The token type to be used for query.
   *
   * Defaults to `TokenType.Access`.
   */
  readonly tokenType?: TokenType;
  /**
   * A generic client implementation generated by ts-proto.
   * e.g. `UserClientImpl`. If the constructor requires additional params, they
   * can be provided via `usePrpcServiceClient(opts, ...params)`.
   */
  readonly ClientImpl: Constructor<S, [PrpcClient, ...Params]> & {
    readonly DEFAULT_SERVICE: string;
  };
}

/**
 * A pRPC method with optional additional params.
 */
type Method<Req, Res, Params extends unknown[] = []> = (
  req: Req,
  ...params: Params
) => Res;

/**
 * The request object from a pRPC method.
 */
type MethodReq<T> = T extends Method<infer Req, infer _Res> ? Req : never;

type PagedReq<Req> = Req extends { readonly pageToken: string } ? Req : never;
type PagedRes<Res> =
  Awaited<Res> extends { readonly nextPageToken: string } ? Res : never;

type MethodRes<T> = T extends Method<infer _Req, infer Res> ? Res : never;

type MethodParams<T> =
  T extends Method<infer _Req, infer _Res, infer Params> ? Params : never;

type MethodKeys<S, Req, Res, Params extends unknown[] = []> = keyof {
  [MK in keyof S as S[MK] extends Method<Req, Res, Params>
    ? MK
    : never]: unknown;
};

export type DecoratedMethod<MK, Req, Res, Params extends unknown[] = []> = {
  (req: Req, ...params: Params): Res;
  /**
   * Builds a ReactQuery option that queries the method. The query key is
   * consisted of
   * `[userIdentity, 'prpc', serviceHost, serviceName, additionalHeaders, methodName, request]`.
   */
  query(
    req: Req,
    ...params: Params
  ): {
    queryKey: [
      string,
      'prpc',
      string,
      string,
      Readonly<Record<string, string>>,
      MK,
      Req,
    ];
    queryFn: () => Res;
  };
  /**
   * Builds a ReactQuery option that queries the paginated method. The next page
   * param is automatically extracted from the previous response. The query key
   * is consisted of
   * `[userIdentity, 'prpc-paged', serviceHost, serviceName, additionalHeaders, methodName, request]`.
   */
  queryPaged(
    req: PagedRes<Res> extends Res ? PagedReq<Req> : never,
    ...params: Params
  ): {
    queryKey: [
      string,
      'prpc-paged',
      string,
      string,
      Readonly<Record<string, string>>,
      MK,
      Req,
    ];
    queryFn: (ctx: { pageParam?: string }) => Res;
    getNextPageParam: (lastRes: Awaited<Res>) => string | null;
    initialPageParam: '';
  };
};

export type DecoratedClient<S> = {
  // The request type has to be `any` because the argument type must be contra-
  // variant when sub-typing a function.
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  [MK in keyof S]: S[MK] extends Method<any, any>
    ? DecoratedMethod<
        MK,
        MethodReq<S[MK]>,
        MethodRes<S[MK]>,
        MethodParams<S[MK]>
      >
    : S[MK];
};

/**
 * Construct a decorated pRPC client with the users credentials. Add helper
 * functions to each pRPC method to help generating ReactQuery options.
 *
 * Example:
 * ```typescript
 * // Constructs a client.
 * const client = usePrpcServiceClient({
 *   // The host of the pRPC server.
 *   host: 'cr-buildbucket-dev.appspot.com',
 *   // The client implementation generated by ts-proto.
 *   ClientImpl: BuildsClientImpl,
 * })
 *
 * // Use the `.query` helper function to build a ReactQuery.
 * const {data, isLoading, ...} = useQuery({
 *   // `.query` is available on all pRPC methods. It takes the same parameters
 *   // as the method itself, and returns a ReactQuery options object with
 *   // `queryKey` and `queryFn` populated.
 *   ...client.GetBuild.query({
 *     // A type checked request object. The type of the request object is
 *     // inferred from the supplied client implementation and method name.
 *     id: "1234",
 *     ...
 *   }),
 *   // Add or override query options if needed.
 *   select: (res) => {...},
 *   ...
 * })
 *
 * // Use the `.queryPaged` helper function to build an infinite ReactQuery.
 * const {data, isLoading, ...} = useInfiniteQuery({
 *   // `.queryPaged` similar to `.query` except that
 *   // 1. it is only available on pRPC methods that have 'pageToken' in the
 *   // request object and 'nextPageToken' in the response object, and
 *   // 2. in addition to `queryKey` and `queryFn`, it also populates
 *   // `getNextPageParam`, making it suitable for `useInfiniteQuery`.
 *   ...client.SearchBuilds.queryPaged({
 *     // A type checked request object. The type of the request object is
 *     // inferred from the supplied client implementation and method name.
 *     id: "1234",
 *     ...
 *   }),
 *   // Add or override query options if needed.
 *   select: (res) => {...},
 *   ...
 * })
 * ```
 */
export function usePrpcServiceClient<
  S extends object,
  Params extends unknown[],
>(opts: PrpcServiceClientOptions<S, Params>, ...serializableParams: Params) {
  const {
    host,
    insecure = false,
    additionalHeaders,
    ClientImpl,
    tokenType = TokenType.Access,
  } = opts;

  const { identity } = useAuthState();
  const getAuthToken = useGetAuthToken(tokenType);
  const additionalHeadersObj = Object.fromEntries(
    new Headers(additionalHeaders).entries(),
  );

  const client = useSingleton({
    key: [
      // Include package name in the key to avoid collision just in case we ends
      // up using the same key format in other places.
      '@/common/hooks/prpc_query',
      getObjectId(ClientImpl),
      host,
      insecure,
      getObjectId(getAuthToken),
      additionalHeaders,
      serializableParams,
    ],
    fn: () =>
      new ClientImpl(
        new PrpcClient({ host, insecure, getAuthToken, additionalHeaders }),
        ...serializableParams,
      ),
  });

  return new Proxy(client, {
    get(target, p, receiver) {
      const mk = p as MethodKeys<S, unknown, unknown>;
      const value = Reflect.get(target, mk, receiver);
      if (typeof value !== 'function') {
        return value;
      }
      const fn = (...args: unknown[]) => value.apply(target, args);
      fn.query = (req: object, ...params: unknown[]) => ({
        queryKey: [
          identity,
          'prpc',
          host,
          ClientImpl.DEFAULT_SERVICE,
          additionalHeadersObj,
          mk,
          req,
        ],
        queryFn: () => fn(req, ...params),
      });
      fn.queryPaged = (req: object, ...params: unknown[]) => ({
        queryKey: [
          identity,
          'prpc-paged',
          host,
          ClientImpl.DEFAULT_SERVICE,
          additionalHeadersObj,
          mk,
          req,
        ],
        queryFn: ({ pageParam = '' }) =>
          fn(pageParam ? { ...req, pageToken: pageParam } : req, ...params),
        // Return `null` when the next page token is an empty string so it get
        // treated as no next page.
        getNextPageParam: ({ nextPageToken = '' }) => nextPageToken || null,
      });
      return fn;
    },
  }) as DecoratedClient<S>;
}
